package com.eshore.textanalyzer.elk.utils;

import com.eshore.textanalyzer.elk.common.EsState;
import com.eshore.textanalyzer.elk.pojo.TagEsState;
import com.eshore.textanalyzer.elk.service.TagEsStateService;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.bulk.BackoffPolicy;
import org.elasticsearch.action.bulk.BulkProcessor;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.unit.TimeValue;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;

/**
 * @author xyd
 * @date 2021/5/13
 * 创建bulkProcessor并初始化
 */
@Slf4j
@Component
public class BulkProcessorUtil {
    @Autowired
    private TagEsStateService tagEsStateService;

    public BulkProcessor getBulkProcessor(RestHighLevelClient client, TagEsState tagEsState) {
        BulkProcessor bulkProcessor = null;

        try {

            BulkProcessor.Listener listener = new BulkProcessor.Listener() {
                /**
                 * 批量执行之前 过程
                 * @param executionId
                 * @param request
                 */
                @Override
                public void beforeBulk(long executionId, BulkRequest request) {
                    log.info("TRY to insert data number: " + request.numberOfActions());
                }

                /**
                 * 批量执行之后 成功
                 * 这里对执行成功的数据入库
                 * @param executionId
                 * @param request
                 * @param response
                 */
                @Override
                public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {
                    log.info("********* Success insert data number ： " + request.numberOfActions() + ", executionId：" + executionId);
                }

                /**
                 * 批量新增后 出错
                 * 在这里对出错后的信息入库
                 * @param executionId
                 * @param request
                 * @param failure
                 */
                @Override
                public void afterBulk(long executionId, BulkRequest request, Throwable failure) {
                    System.out.println(failure.getMessage());
                    log.error("Bulk is unsuccess ： " + failure + ", executionId：" + executionId);
                }
            };

            BiConsumer<BulkRequest, ActionListener<BulkResponse>> bulkConsumer = (request, bulkListener) -> client.bulkAsync(request, RequestOptions.DEFAULT,bulkListener);

            BulkProcessor.Builder builder = BulkProcessor.builder(bulkConsumer,listener);
            // 设置何时刷新批量请求

            // 服务器上50000条每批或者10000条每批 本地上5000就好
            builder.setBulkActions(10000);
            // 服务器上设置300L 本地100L就好
            builder.setBulkSize(new ByteSizeValue(100L, ByteSizeUnit.MB));
            builder.setConcurrentRequests(10);
            builder.setFlushInterval(TimeValue.timeValueSeconds(100L));
            builder.setBackoffPolicy(BackoffPolicy.constantBackoff(TimeValue.timeValueSeconds(1L), 3));

            bulkProcessor = builder.build();

        }catch (Exception e) {
            String end_time = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").format(new Date());
            tagEsState.setEndTime(end_time);
            tagEsState.setExecutingState(EsState.FAILED);
            tagEsState.setMsg(e.getMessage());
            tagEsStateService.updateState(tagEsState);
            e.printStackTrace();
            try {
                boolean terminatedFlag = bulkProcessor.awaitClose(150L, TimeUnit.SECONDS);
                client.close();
                log.info(String.valueOf(terminatedFlag));
            } catch (Exception e1) {
                log.error(e1.getMessage());
            }
        }
        return bulkProcessor;
    }
}
